#!/usr/bin/env python
# -*- coding: UTF-8 -*-
# @Project : mcp-gateway
# @File    : app.py
# @IDE     : PyCharm
# @Author  :ZH
# @Time    : 2025/6/20 18:10
import asyncio
import os
import subprocess
import uuid
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from typing import Dict
import socket

app = FastAPI()

# 存储运行中的任务 {task_id: process}
tasks: Dict[str, subprocess.Popen] = {}
tasks_port: Dict[str, int] = {}

# Nginx 配置路径
NGINX_CONF_DIR = "proxy.d"
NGINX_CMD_RELOAD = ["nginx", "-s", "reload"]


class CommandRequest(BaseModel):
    cmd: str


def find_free_port():
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
        s.bind(("", 0))
        return s.getsockname()[1]


def update_nginx_config(task_id: str, port: int):
    """
    生成并写入 Nginx 配置文件，优化 SSE 支持
    """
    config_content = f"""
# 针对根路径的请求
location /{task_id} {{
    rewrite ^/{task_id}$ /{task_id}/ redirect;
}}

location /{task_id}/ {{
    proxy_pass http://127.0.0.1:{port}/;
    proxy_http_version 1.1;
    proxy_set_header Upgrade $http_upgrade;
    proxy_set_header Connection 'upgrade';
    proxy_set_header Host $host;
    proxy_set_header X-Real-IP $remote_addr;
    proxy_cache off;
    proxy_buffering off;
    proxy_read_timeout 86400; # 延长超时时间，适合长连接
}}


# 任务特定的SSE端点
location /{task_id}/sse {{
    proxy_pass http://127.0.0.1:{port}/{task_id}/sse;
    proxy_http_version 1.1;
    proxy_set_header Connection '';
    proxy_set_header Host $host;
    proxy_set_header X-Real-IP $remote_addr;
    proxy_buffering off;
    proxy_cache off;
    proxy_read_timeout 86400s;
    chunked_transfer_encoding on;
    proxy_set_header X-Forwarded-For $proxy_add_x_forwarded_for;
}}


location /{task_id}/message {{
    proxy_pass http://127.0.0.1:{port}/{task_id}/message;
    proxy_http_version 1.1;
    proxy_set_header Upgrade $http_upgrade;
    proxy_set_header Connection 'upgrade';
    proxy_set_header Host $host;
    proxy_set_header X-Real-IP $remote_addr;
    proxy_cache off;
    proxy_buffering off;
}}
"""

    # 确保配置目录存在
    os.makedirs(NGINX_CONF_DIR, exist_ok=True)

    config_path = os.path.join(NGINX_CONF_DIR, f"task_{task_id}.conf")
    with open(config_path, "w") as f:
        f.write(config_content)

    # 重载 Nginx
    try:
        subprocess.run(NGINX_CMD_RELOAD, check=True)
    except subprocess.CalledProcessError as e:
        raise RuntimeError(f"Nginx reload failed: {e}")


def remove_nginx_config(task_id: str):
    """
    删除 Nginx 配置文件
    """
    config_path = os.path.join(NGINX_CONF_DIR, f"task_{task_id}.conf")
    if os.path.exists(config_path):
        os.remove(config_path)

    # 重载 Nginx
    try:
        subprocess.run(NGINX_CMD_RELOAD, check=True)
    except subprocess.CalledProcessError as e:
        raise RuntimeError(f"Nginx reload failed: {e}")


@app.post("/start")
async def start_command(command: CommandRequest) -> dict:
    task_id = str(uuid.uuid4())

    try:
        port = find_free_port()
        full_cmd = f'supergateway --stdio "{command.cmd} /mcp_server/tmp/" --port {port} --ssePath /{task_id}/sse --messagePath /{task_id}/message'

        process = subprocess.Popen(
            full_cmd,
            shell=True,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            stdin=subprocess.PIPE
        )

        tasks[task_id] = process
        tasks_port[task_id] = port

        # 更新 Nginx 配置
        update_nginx_config(task_id, port)
        # 需要等待服务完成，等待10秒返回
        await asyncio.sleep(10)

        return {"task_id": task_id, "status": "started"}

    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@app.post("/stop/{task_id}")
async def stop_command(task_id: str) -> dict:
    process = tasks.get(task_id)
    if not process:
        raise HTTPException(status_code=404, detail="Task not found")

    try:
        process.terminate()
        try:
            process.wait(timeout=3)
        except subprocess.TimeoutExpired:
            process.kill()

        del tasks[task_id]
        del tasks_port[task_id]

        # 删除 Nginx 配置
        remove_nginx_config(task_id)

        return {"task_id": task_id, "status": "stopped"}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


if __name__ == '__main__':
    import uvicorn

    uvicorn.run('app:app', port=8000, host='0.0.0.0')
